# 主要目标: 按月度分析哪个月应该买入，哪个月应该卖出股票数据
# 日志
# 2025/01/28 新建本脚本，并初步实现基本功能。
# 2025/02/02 读取日K信息，生成月K，再生成二维表，最后生成统计数据
# 2025/02/03 
import baostock as bs
import pandas as pd
from datetime import datetime
import os,sys
import matplotlib.pyplot as plt
import seaborn as sns

# 创建一个函数，用于统计正数和负数的个数
def count_positive_negative(values):
    positive_count = (values > 0).sum()
    negative_count = (values < 0).sum()
    return positive_count, negative_count

def get_csvfiles(directory)
    # 获取目录下所有CSV文件全路径
    paths = []
    try:
        # 遍历目录
        for filename in os.listdir(directory):
            full_path = os.path.join(directory, filename).lower()
            if os.path.isfile(full_path) and full_path.endswith('.csv'):
                paths.append(full_path)
    except FileNotFoundError:
        print(f"目录 {directory} 不存在")
    except PermissionError:
        print(f"没有权限访问目录 {directory}")
    return paths

# 指定目录路径
directory = sys.argv[1] if len(sys.argv) > 1 else r'c:\gitee.com\stock_research\samples'
print("Target Directory:", directory)
id = 0

def process_file(csvfile):
    ps = []
    # 排除各类指数
    if 'sh.0' in csvfile or 'sz.399' in csvfile:
        return ps
    
    # 如果不是sh.或sz.开头则不是目标
    if not ('sh.' in csvfile or 'sz.' in csvfile):
        return ps

    # 1. 加载数据
    df = pd.read_csv(csvfile, index_col=0)

    # 2.将日期列转换为datetime格式并设置为索引，将其他值改为浮点值。
    df['date'] = pd.to_datetime(df['date'])
    df.set_index('date', inplace=True)
    df[['open', 'high', 'low', 'close']] = df[['open', 'high', 'low', 'close']].astype('float64')  # 将其他列转换为浮点型

    # 3. 按年和月分组
    grouped = df.groupby([df.index.year, df.index.month])

    # 将日K转为月K，并保存在 monthly_data 列表中。
    monthly_data = []
    # 遍历分组，计算月K数据
    for (year, month), group in grouped:
        # 获取当月第一天和最后一天的数据
        first_day = group.index.min()
        last_day = group.index.max()
        # 获取当月第一天的开盘价
        monthly_open = group.loc[first_day, 'open']
        # 获取当月最后一天的收盘价
        monthly_close = group.loc[last_day, 'close']
        # 获取当月最高价和最低价
        monthly_high = group['high'].max()
        monthly_low = group['low'].min()
        # 获取上个月的收盘价
        keys=list(grouped.groups.keys())
        if (year, month) == (keys[0][0], keys[0][1]):
            # 如果是第一组数据，没有上个月的收盘价，可以设置为NaN或某个初始值
            monthly_preclose = None
        else:
            # 获取上个月最后一天的收盘价
            previous_month_last_day = grouped.get_group((year, month)).index.max()
            monthly_preclose = grouped.get_group((year, month)).loc[previous_month_last_day, 'close']
        # 将计算结果添加到列表中
        monthly_data.append({
            'year': year,
            'month': month,
            'preclose': monthly_preclose,
            'open': monthly_open,
            'high': monthly_high,
            'low': monthly_low,
            'close': monthly_close
        })

    # 将月K数据列表转换为DataFrame
    monthly_df = pd.DataFrame(monthly_data)
    monthly_df['increment'] = monthly_df['close'] - monthly_df['open']
    # 创建一个以年为纵坐标，月为横坐标的表格
    pivot_table = monthly_df.pivot(index='year', columns='month', values='increment')
    
    # 只统计满足10年的数据
    if len(pivot_table) < 10:
        return ps
    
    # 按月份分组，并统计每个月的正数和负数个数
    result = monthly_df.groupby('month')['increment'].apply(count_positive_negative).reset_index()
    # 将结果拆分为两列
    result[['positive_count', 'negative_count']] = result['increment'].apply(pd.Series)
    # 删除多余的列
    result = result.drop(columns=['increment'])
    # 打印结果
    sum1 = []
    for index, row in result.iterrows():
        sum1.append((row.tolist()))
        
    
    for vs in sum1:
        ps.append(vs[1]/(vs[1]+vs[2]))

    if max(ps) < 0.99:
        return ps

    # 打印表头（月份）
    print("-" * 124)
    print("Year |", " | ".join(f"{month:^7}" for month in pivot_table.columns))
    print("-" * 124)
    # 打印每一年的数据
    for year, row in pivot_table.iterrows():
        # 格式化每一行的数据
        row_values = [f"{value:7.2f}" if not pd.isna(value) else "  N/A  " for value in row]
        print(f"{year:4} |", " | ".join(row_values))
    print("-" * 124)
    print(' s1 ', end='')
    for p in ps:
        print(f' | {p:7.1%}', end='')
    print('\n' + "-" * 124)
    return ps


paths = get_csvfiles(directory)
total_count = len(paths)
res = []
for id, path in enumerate(paths):
    print(f'id={id:0>4}/{total_count}  {datetime.now():%Y-%m-%d %H:%M:%S} {path}')
    try:
        ps = process_file(path)
        res.append(ps)
    except Exception as e1:
        print(e1)

for p in res:
    print(p)
print('total:', len(res))